import pandas as pd
import numpy as np
from typing import Optional



def mark_col_growth(
    df: pd.DataFrame,
    target_col: str,
    window: int = 10,
    feature_suffix: Optional[str] = None,
    min_valid_value: float = 1e-6,
    growth_type: str = 'relative',
    trading_col: str = 'is_trading',
    is_drop: bool = True
) -> pd.DataFrame:
    """
    优化版增长率计算 - 保持原有4列输出
    
    改进点：
    1. 零值特殊处理：当基准值为零时，使用替代策略计算增长率
    2. 仅保留最终需要的4列输出
    3. 使用临时DataFrame避免污染原始数据
    """
    # 输入校验
    required_cols = {'code', 'date', target_col}
    missing_cols = required_cols - set(df.columns)
    if missing_cols:
        raise ValueError(f"缺失必要列: {missing_cols}")
    
    # 创建临时工作DataFrame（仅包含必要列）
    work_df = df[['code', 'date', target_col]].copy()
    
    # 数据预处理
    work_df['date'] = pd.to_datetime(work_df['date'])
    work_df.sort_values(by=['code', 'date'], inplace=True)
    
    # 特征列命名
    suffix = f"_{feature_suffix}" if feature_suffix else ''
    growth_col = f"{target_col}_yoy{suffix}_{window}d"
    pre_growth_col = f"pre_{growth_col}"
    
    # 零值特殊处理 - 添加微小常数避免除零错误
    epsilon = 1e-10
    work_df['_adjusted_target'] = work_df[target_col].replace(0, epsilon) + epsilon
    
    # 滚动基准计算（优化零值处理）
    work_df['_base'] = work_df.groupby('code', group_keys=False)['_adjusted_target'].transform(
        lambda x: x.shift(1).rolling(window, min_periods=1).mean()
    )
    # display(work_df[-30:])
    # 增长率计算（优化零值处理）
    prev_col = work_df.groupby('code')[target_col].shift(1)
    
    if growth_type == 'relative':
        # 核心优化：零值特殊处理
        with np.errstate(divide='ignore', invalid='ignore'):
            work_df[growth_col] = np.where(
                work_df['_base'] <= min_valid_value,
                # 基准接近零时的替代计算
                (work_df[target_col] - prev_col) / (prev_col + epsilon),
                # 正常相对增长率计算
                (work_df[target_col] - work_df['_base']) / work_df['_base']
            )
        # display(work_df[-30:])
    elif growth_type == 'log':
        # 对数增长率同样需要零值处理
        work_df[growth_col] = np.where(
            (work_df['_base'] > min_valid_value) & (work_df[target_col] > min_valid_value),
            np.log(work_df[target_col] / work_df['_base']),
            np.where(
                # 零到非零的跃迁
                (work_df[target_col] > min_valid_value) & (prev_col <= min_valid_value),
                1.0,  # 固定值表示显著增长
                np.nan
            )
        )
    else:
        raise ValueError("growth_type 必须为 'relative' 或 'log'")
    
    # 计算昨日特征
    work_df[pre_growth_col] = work_df.groupby('code')[growth_col].shift(1)
    work_df[growth_col+'_2'] = work_df['_adjusted_target']/work_df['_base']
    
    # 仅保留最终需要的4列
    result = work_df[['code', 'date', growth_col,growth_col+'_2', pre_growth_col]].copy()
    # 停牌数据处理（如果提供停牌列）
    if trading_col in df.columns:
        print('dddd',trading_col)
        # 合并停牌信息
        trading_info = df[['code', 'date', trading_col]].drop_duplicates()
        result = result.merge(trading_info, on=['code', 'date'], how='left')
        result.loc[~result[trading_col], [growth_col, pre_growth_col]] = np.nan
        result.drop(columns=[trading_col], inplace=True)
    
    return result.round(4)  # 提高精度保留小数

    
import pandas as pd
import numpy as np

def mark_vol_turnover(df):
    """生成预测三日涨超15%的核心特征"""
    df = df.sort_values(['code','date']).copy()
    
    # === 特征工程 ===
  
    # 2. 成交额突破信号（网页3突破模型）
    df['turnover_break'] = (df['turnover'] > df.groupby('code')['turnover']
                            .transform(lambda x: x.rolling(5).max().shift())).astype(int)
    
    # 3. 换手率波动强度（网页7标准化处理）
    df['t_rate_zscore'] = df.groupby('code')['t_rate'].transform(
        lambda x: (x - x.rolling(10).mean()) / x.rolling(10).std().replace(0, 1e-6))
    
    
    # # 5. 量能效率指标（网页2趋势模型） TODO 使用了未来数据
    # df['vol_eff'] = df.groupby('code')['close'].transform(
    #     lambda x: x.pct_change(3).shift(-3)) / df['volume'].rolling(3).sum()
    
    return df[['code','date','turnover_break',
              't_rate_zscore']]
